"""
督导系统模型配置文件
包含风险预警、问题分类、热力图预测等模型的配置参数
"""

import os
from dataclasses import dataclass
from typing import Dict, List, Optional

@dataclass
class RiskPredictionConfig:
    """风险预警模型配置"""
    model_name: str = "risk_prediction_model"
    model_type: str = "xgboost"  # xgboost, lightgbm, catboost
    
    # 特征工程
    feature_window_days: int = 30  # 历史数据窗口期
    spatial_grid_size: float = 0.01  # 空间网格大小(度)
    temporal_aggregation: str = "daily"  # daily, weekly, monthly
    
    # 模型参数
    n_estimators: int = 1000
    max_depth: int = 8
    learning_rate: float = 0.1
    random_state: int = 42
    
    # 训练参数
    test_size: float = 0.2
    val_size: float = 0.1
    cv_folds: int = 5
    early_stopping_rounds: int = 100

@dataclass 
class ProblemClassificationConfig:
    """问题分类模型配置"""
    model_name: str = "problem_classification_model"
    base_model: str = "damo/nlp_structbert_backbone_base_std"  # 魔搭社区中文BERT模型
    
    # 文本处理
    max_length: int = 512
    text_fields: List[str] = None
    
    # 分类标签
    problem_types: List[str] = None
    severity_levels: List[str] = None
    
    # 训练参数
    batch_size: int = 16
    learning_rate: float = 2e-5
    num_epochs: int = 10
    warmup_steps: int = 500
    weight_decay: float = 0.01
    
    def __post_init__(self):
        if self.text_fields is None:
            self.text_fields = ["问题描述", "现场记录", "备注"]
        
        if self.problem_types is None:
            self.problem_types = [
                "安全隐患", "环境污染", "违规施工", "质量问题", 
                "制度执行", "人员管理", "设备故障", "其他"
            ]
            
        if self.severity_levels is None:
            self.severity_levels = ["低", "中", "高", "紧急"]

@dataclass
class HeatmapPredictionConfig:
    """热力图预测模型配置"""
    model_name: str = "heatmap_prediction_model"
    model_type: str = "spatial_temporal_cnn"
    
    # 空间配置
    grid_resolution: float = 0.005  # 网格分辨率
    spatial_window_km: float = 5.0  # 空间影响范围
    
    # 时间配置  
    temporal_window_hours: int = 168  # 时间窗口(小时)
    prediction_horizon_hours: int = 24  # 预测时间范围
    
    # 模型结构
    spatial_filters: List[int] = None
    temporal_filters: List[int] = None
    hidden_dims: List[int] = None
    dropout_rate: float = 0.3
    
    # 训练参数
    batch_size: int = 32
    learning_rate: float = 0.001
    num_epochs: int = 50
    patience: int = 10
    
    def __post_init__(self):
        if self.spatial_filters is None:
            self.spatial_filters = [64, 128, 256]
        if self.temporal_filters is None:
            self.temporal_filters = [32, 64, 128]
        if self.hidden_dims is None:
            self.hidden_dims = [512, 256, 128]

@dataclass
class DataConfig:
    """数据配置"""
    # 数据库连接
    db_host: str = os.getenv("DB_HOST", "localhost")
    db_port: int = int(os.getenv("DB_PORT", "5432"))
    db_name: str = os.getenv("DB_NAME", "supervision_db")
    db_user: str = os.getenv("DB_USER", "postgres")
    db_password: str = os.getenv("DB_PASSWORD", "password")
    
    # 数据路径
    data_dir: str = "data"
    model_dir: str = "models"
    log_dir: str = "logs"
    
    # 数据处理
    min_samples_per_area: int = 10
    outlier_threshold: float = 3.0
    missing_value_threshold: float = 0.3

class ModelConfig:
    """主配置类"""
    def __init__(self):
        self.risk_prediction = RiskPredictionConfig()
        self.problem_classification = ProblemClassificationConfig()
        self.heatmap_prediction = HeatmapPredictionConfig()
        self.data = DataConfig()
        
        # 创建必要目录
        for dir_path in [self.data.data_dir, self.data.model_dir, self.data.log_dir]:
            os.makedirs(dir_path, exist_ok=True)

# 全局配置实例
config = ModelConfig() 